package com.superzlc.utils;

import java.util.HashMap;

import javax.crypto.Mac;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;

import org.apache.commons.codec.binary.Base64;
import org.apache.commons.codec.binary.Hex;

public abstract class OTP {

	private static final String MAC_NAME = "HmacSHA1";

	public static void main(String[] args) {

		HOTP p1 = new HOTP("123456", 0);
		System.out.println(p1.gen()); // 186818
		System.out.println();

		System.out.println(p1.gen(1234)); // 263197
		System.out.println(p1.verify("263197", 1234));
		System.out.println();

		TOTP p2 = new TOTP("123456", 3, 30);
		System.out.println(p2.gen(1614054214456L)); // 330925
		System.out.println(p2.verify("330925", 1614054214456L));

		TOTP p3 = new TOTP("helloworld", 3, 30);
		p3.setSecret("helloworld", "string");
		p3.setSecret("68656c6c6f776f726c64", "hex");
		p3.setSecret("aGVsbG93b3JsZA==", "base64");
		System.out.println(p3.gen(1614054214456L)); // 759440
		System.out.println(p3.verify("759440", 1614054214456L));

	}

	// ----------------
	public static class HOTP {
		private byte[] secret = null;
		private long counter = 0;

		public HOTP() {
			this.secret = null;
			this.counter = 0;
		}

		public HOTP(String secret, long counter) {
			this.secret = secret == null ? null : secret.getBytes();
			this.counter = counter;
		}

		public void setSecret(byte[] secret) {
			this.secret = secret;
		}

		public void setSecret(String secret, String format) {
			this.secret = _convertSecret(secret, format);
		}

		//
		public String gen() {
			return gen(this.counter);
		}

		public boolean verify(String token) {
			if (verify(token, this.counter)) {
				this.counter++;
				return true;
			}
			return false;
		}

		public String gen(long counter) {
			return _gen(secret, counter);
		}

		public boolean verify(String token, long counter) {
			return token.equals(gen(counter));
		}

		// ------------
		private static String _gen(byte[] secret, long counter) {
			byte[] data = new byte[8];
			for (int i = 7; i >= 0; i--) {
				data[i] = (byte) (counter & 0xFF);
				counter >>>= 8;
			}
			byte[] hash = _hmacSHA1Encrypt(secret, data);
			int offset = hash[19] & 0xf;
			int num = (hash[offset] & 0x7f) << 24 |
				(hash[offset + 1] & 0xff) << 16 |
				(hash[offset + 2] & 0xff) << 8  |
				(hash[offset + 3] & 0xff);
			num = num % 1000000;
			return String.format("%06d", num);
		}

		//
		private static byte[] _hmacSHA1Encrypt(byte[] secret, byte[] data) {
			try {
				Mac mac = Mac.getInstance(MAC_NAME);
				SecretKey secretKey = new SecretKeySpec(secret, MAC_NAME);
				mac.init(secretKey);
				return mac.doFinal(data);
			} catch(Exception e) {
				throw new RuntimeException("hmacSHA1加密失败", e);
			}
		}

		//
		private static byte[] _convertSecret(String secret, String format) {
			if ("hex".equals(format)) {
				try {
					return Hex.decodeHex(secret);
				} catch(Exception e) {
					throw new RuntimeException("参数错误");
				}
			} else if ("base64".equals(format)) {
				return Base64.decodeBase64(secret);
			} else if ("string".equals(format)) {
				return secret.getBytes();
			} else {
				throw new RuntimeException("参数错误");
			}
		}

	}

	public static class TOTP {
		private HOTP hotp = null;
		private int window = 3;
		private int interval = 30;

		public TOTP() {
			this.hotp = new HOTP();
			this.window = 3;
			this.interval = 30;
		}

		public TOTP(String secret, int window, int interval) {
			this.hotp = new HOTP(secret, 0);
			this.window = window;
			this.interval = interval;
		}

		public void setSecret(byte[] secret) {
			this.hotp.setSecret(secret);
		}

		public void setSecret(String secret, String format) {
			this.hotp.setSecret(secret, format);
		}

		//
		public String gen() {
			long timestamp = System.currentTimeMillis();
			return gen(timestamp);
		}

		public String gen(long timestamp) {
			long counter = (long) (timestamp / 1000 / this.interval);
			return this.hotp.gen(counter);
		}

		public HashMap<String, Object> gen2() {
			long timestamp = System.currentTimeMillis();
			return gen2(timestamp);
		}

		public HashMap<String, Object> gen2(long timestamp) {
			long timestamp_second = timestamp / 1000;
			long counter = (long) (timestamp_second / this.interval);
			String token = this.hotp.gen(counter);
			long step = timestamp_second % this.interval;
			long refresh_time = counter * this.interval;

			HashMap<String, Object> map = new HashMap<>();
			map.put("token", token);
			map.put("loop_length", this.interval);
			map.put("loop_step", step);
			map.put("refresh_interval", this.interval * 1000);
			map.put("refresh_time", refresh_time * 1000);
			map.put("timestamp", timestamp);
			return map;
		}

		public boolean verify(String token) {
			return this.verify(token, System.currentTimeMillis());
		}

		public boolean verify(String token, long timestamp) {
			long counter = (long) (timestamp / 1000 / interval);

			for (int i = 0, w = 0; ; i++) {
				if (token.equals(this.hotp.gen(counter - i)))
					return true;
				if (++w >= window)
					break;
				if (i > 0) {
					if (token.equals(this.hotp.gen(counter + i)))
						return true;
					if (++w > window)
						break;
				}
			}
			return false;
		}

	}

}
